import csv

def clean_student_faculty_ratio(ratio_str):
    """Clean and extract the numeric value from student-faculty ratio string."""
    if not ratio_str:
        return None
    # Handle different ratio formats seen in the data
    ratio_str = ratio_str.replace("：", ":")  # Replace Chinese colon if present
    if ":" in ratio_str:
        try:
            parts = ratio_str.split(":")
            # Some entries have extra text before the ratio (e.g., "0112.4")
            if len(parts) > 1:
                return float(parts[1].strip())
            else:
                return None
        except ValueError:
            return None
    return None

def process_school_data(input_file, output_file):
    data = []
    
    with open(input_file, mode="r", encoding="utf-8-sig") as file:
        reader = csv.reader(file)
        for row in reader:
            # Skip empty rows
            if not row or not row[0]:
                continue
            data.append(row)
    
    # Standardize headers (assuming first row is header)
    headers = data[0]
    header_map = {
        "序号": "serial_no",
        "学校名称": "school_name",
        "学校性质": "school_type",
        "教职工人数": "total_staff_count",
        "专职教师人数": "full_time_teacher_count",
        "师生比": "student_faculty_ratio_raw"
    }
    
    # Rename headers
    for i in range(len(headers)):
        if headers[i] in header_map:
            headers[i] = header_map[headers[i]]
    
    # Filter and process data
    cleaned_data = []
    cleaned_data.append([
        "serial_no",
        "school_name",
        "total_staff_count",
        "full_time_teacher_count",
        "estimated_student_count",
        "students_per_ft_teacher"
    ])
    
    for row in data[1:]:  # Skip header row
        if len(row) < len(headers):
            continue  # Skip incomplete rows
            
        school_type = row[headers.index("school_type")] if "school_type" in headers else ""
        
        # Only process public schools
        if school_type != "公办":
            continue
            
        try:
            # Get required fields
            serial_no = int(row[headers.index("serial_no")])
            school_name = row[headers.index("school_name")]
            total_staff = int(row[headers.index("total_staff_count")])
            full_time_teachers = int(row[headers.index("full_time_teacher_count")])
            ratio_str = row[headers.index("student_faculty_ratio_raw")]
            
            # Process ratio
            ratio_value = clean_student_faculty_ratio(ratio_str)
            if ratio_value is None:
                continue  # Skip if ratio couldn't be parsed
                
            # Calculate estimated student count
            estimated_students = round(total_staff * ratio_value)
            
            # Calculate students per full-time teacher (handle division by zero)
            students_per_teacher = round(estimated_students / full_time_teachers, 1) if full_time_teachers > 0 else None
            
            if students_per_teacher is None:
                continue  # Skip if division by zero would occur
                
            cleaned_data.append([
                serial_no,
                school_name,
                total_staff,
                full_time_teachers,
                estimated_students,
                students_per_teacher
            ])
            
        except (ValueError, IndexError):
            continue  # Skip rows with invalid data
    
    # Save cleaned data
    with open(output_file, mode="w", encoding="utf-8", newline="") as file:
        writer = csv.writer(file)
        writer.writerows(cleaned_data)

# Process the 2020 data
process_school_data("baoshan-schools-2020.csv", "clean-baoshan-schools-2020.csv")